package org.budo.elasticsearch.jdbc.driver.statement;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcConnection;
import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcPreparedStatement;
import org.budo.elasticsearch.jdbc.driver.parser.jsql.JSqlParserStatement;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.client.transport.TransportClient;

import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.insert.Insert;

/**
 * @author limingwei
 */
public class InsertStatement extends AbstractElasticsearchStatement {
    private static final String ID_KEY = "_id";

    @Override
    public void execute(PreparedStatement preparedStatement) throws SQLException {
        List<Map<String, Object>> insertSourceList = this.insertSource(preparedStatement);
        if (insertSourceList.size() == 1) { // 单条 insert
            this.index(preparedStatement, insertSourceList.get(0));
            return;
        }

        this.bulk(preparedStatement, insertSourceList); // 批量insert
    }

    private void bulk(PreparedStatement preparedStatement, List<Map<String, Object>> insertSourceList) throws SQLException {
        BudoElasticsearchJdbcConnection connection = (BudoElasticsearchJdbcConnection) preparedStatement.getConnection();
        TransportClient transportClient = connection.getTransportClient();

        BulkRequestBuilder bulkRequest = transportClient.prepareBulk();
        for (Map<String, Object> insertSource : insertSourceList) {
            IndexRequestBuilder indexRequest = this.indexRequest(connection, insertSource);
            bulkRequest.add(indexRequest);
        }

        this.executeAndGet(bulkRequest);

        ((BudoElasticsearchJdbcPreparedStatement) preparedStatement).setUpdateCount(insertSourceList.size());
    }

    private void index(PreparedStatement preparedStatement, Map<String, Object> insertSource) throws SQLException {
        BudoElasticsearchJdbcConnection connection = (BudoElasticsearchJdbcConnection) preparedStatement.getConnection();

        IndexRequestBuilder indexRequest = this.indexRequest(connection, insertSource);

        this.executeAndGet(indexRequest);

        ((BudoElasticsearchJdbcPreparedStatement) preparedStatement).setUpdateCount(1);
    }

    private IndexRequestBuilder indexRequest(BudoElasticsearchJdbcConnection connection, Map<String, Object> insertSource) throws SQLException {
        TransportClient transportClient = connection.getTransportClient();

        IndexRequestBuilder indexRequest = transportClient.prepareIndex() //
                .setIndex(connection.getCatalog()) //
                .setType(this.tableName());

        Object idValue = insertSource.get(ID_KEY);

        if (null != idValue && !idValue.toString().isEmpty()) {
            indexRequest.setId(idValue.toString());
            insertSource.remove(ID_KEY);
        }

        indexRequest.setSource(insertSource);
        return indexRequest;
    }

    private List<Map<String, Object>> insertSource(PreparedStatement preparedStatement) {
        ItemsList itemsList = ((Insert) ((JSqlParserStatement) this.getSqlStatement()).getStatement()).getItemsList();
        List<Map<String, Object>> list = new ArrayList<Map<String, Object>>();

        if (itemsList instanceof MultiExpressionList) {
            MultiExpressionList multiExpressionList = (MultiExpressionList) itemsList;
            for (ExpressionList expressionList : multiExpressionList.getExprList()) {
                Map<String, Object> _insertSource = this._insertSource(preparedStatement, expressionList);
                list.add(_insertSource);
            }
            return list;
        }

        Map<String, Object> _insertSource = this._insertSource(preparedStatement, (ExpressionList) itemsList);
        list.add(_insertSource);
        return list;
    }

    private Map<String, Object> _insertSource(PreparedStatement preparedStatement, ExpressionList expressionList) {
        List<Expression> expressions = expressionList.getExpressions();

        List<Column> columns = ((Insert) ((JSqlParserStatement) this.getSqlStatement()).getStatement()).getColumns();

        Map<String, Object> source = new HashMap<String, Object>();
        for (int i = 0; i < columns.size(); i++) {
            Column column = columns.get(i);
            String key = this.columnToEsKey(column);

            Expression expression = expressions.get(i);
            Object value = this.expressionValue(preparedStatement, expression);
            source.put(key, value);
        }

        return source;
    }
}